import json
import logging
from datetime import datetime
import pandas as pd
import pdfkit
from jinja2 import Template, FileSystemLoader, Environment
import os
from io import BytesIO
from app.models.database import db, Scan, Vulnerability
from app.services.risk_assessment import get_risk_distribution, get_risk_timeline

logger = logging.getLogger(__name__)

class ReportService:
    """报告服务，负责生成扫描报告和导出结果"""
    
    # 报告模板路径
    TEMPLATE_DIR = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'templates', 'reports')
    
    @staticmethod
    def generate_report(scan_id, format_type='json'):
        """生成扫描报告
        
        Args:
            scan_id (str): 扫描任务ID
            format_type (str): 报告格式，支持 'json', 'html', 'pdf', 'excel'
            
        Returns:
            tuple: (报告数据, 文件扩展名, MIME类型)
        """
        try:
            # 获取扫描数据
            scan = Scan.query.get(scan_id)
            if not scan:
                raise ValueError(f"未找到ID为{scan_id}的扫描任务")
            
            # 获取漏洞数据
            vulnerabilities = Vulnerability.query.filter_by(scan_id=scan_id).all()
            vuln_list = [v.to_dict() for v in vulnerabilities]
            
            # 计算统计信息
            total_vulns = len(vulnerabilities)
            high_vulns = len([v for v in vulnerabilities if v.severity == 'high'])
            medium_vulns = len([v for v in vulnerabilities if v.severity == 'medium'])
            low_vulns = len([v for v in vulnerabilities if v.severity == 'low'])
            
            # 按漏洞类型统计
            vuln_types = {}
            for vuln in vulnerabilities:
                if vuln.vuln_type not in vuln_types:
                    vuln_types[vuln.vuln_type] = 0
                vuln_types[vuln.vuln_type] += 1
            
            # 获取风险分布和时间线数据
            risk_distribution = get_risk_distribution(scan_id)
            risk_timeline = get_risk_timeline(scan_id)
            
            # 准备报告数据
            report_data = {
                'scan_info': {
                    'scan_id': scan.id,
                    'target': scan.target,
                    'status': scan.status,
                    'created_at': scan.created_at.isoformat() if scan.created_at else None,
                    'finished_at': scan.finished_at.isoformat() if scan.finished_at else None
                },
                'statistics': {
                    'total': total_vulns,
                    'high': high_vulns,
                    'medium': medium_vulns,
                    'low': low_vulns,
                    'types': vuln_types
                },
                'risk_metrics': {
                    'risk_distribution': risk_distribution,
                    'risk_timeline': risk_timeline
                },
                'vulnerabilities': vuln_list,
                'generated_at': datetime.utcnow().isoformat()
            }
            
            # 根据格式类型生成不同的报告
            if format_type.lower() == 'json':
                return json.dumps(report_data, ensure_ascii=False, indent=2), 'json', 'application/json'
            elif format_type.lower() == 'html':
                return ReportService._generate_html_report(report_data), 'html', 'text/html'
            elif format_type.lower() == 'pdf':
                return ReportService._generate_pdf_report(report_data), 'pdf', 'application/pdf'
            elif format_type.lower() == 'excel':
                return ReportService._generate_excel_report(report_data), 'xlsx', 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet'
            else:
                raise ValueError(f"不支持的报告格式: {format_type}")
                
        except Exception as e:
            logger.error(f"生成报告失败: {str(e)}")
            raise
    
    @staticmethod
    def _generate_html_report(report_data):
        """生成HTML格式的报告"""
        try:
            # 创建Jinja2环境
            env = Environment(loader=FileSystemLoader(ReportService.TEMPLATE_DIR))
            
            # 尝试加载自定义模板，如果不存在则使用默认模板
            try:
                template = env.get_template('detailed_report.html')
            except FileNotFoundError:
                # 如果没有自定义模板，使用内联的简单模板
                template = Template('''
                <!DOCTYPE html>
                <html lang="zh-CN">
                <head>
                    <meta charset="UTF-8">
                    <meta name="viewport" content="width=device-width, initial-scale=1.0">
                    <title>漏洞扫描报告 - {{ report_data.scan_info.target }}</title>
                    <style>
                        body { font-family: Arial, sans-serif; margin: 20px; }
                        h1, h2 { color: #333; }
                        table { border-collapse: collapse; width: 100%; margin-top: 20px; }
                        th, td { border: 1px solid #ddd; padding: 8px; text-align: left; }
                        th { background-color: #f2f2f2; }
                        .summary { background-color: #f9f9f9; padding: 15px; border-radius: 5px; margin-bottom: 20px; }
                        .severity-high { color: #d32f2f; }
                        .severity-medium { color: #f57c00; }
                        .severity-low { color: #388e3c; }
                    </style>
                </head>
                <body>
                    <h1>漏洞扫描报告</h1>
                    
                    <div class="summary">
                        <h2>扫描信息</h2>
                        <p><strong>目标:</strong> {{ report_data.scan_info.target }}</p>
                        <p><strong>扫描ID:</strong> {{ report_data.scan_info.scan_id }}</p>
                        <p><strong>状态:</strong> {{ report_data.scan_info.status }}</p>
                        <p><strong>开始时间:</strong> {{ report_data.scan_info.created_at }}</p>
                        <p><strong>结束时间:</strong> {{ report_data.scan_info.finished_at }}</p>
                        <p><strong>生成时间:</strong> {{ report_data.generated_at }}</p>
                    </div>
                    
                    <div class="summary">
                        <h2>漏洞统计</h2>
                        <p><strong>总漏洞数:</strong> {{ report_data.statistics.total }}</p>
                        <p><strong>高危漏洞:</strong> <span class="severity-high">{{ report_data.statistics.high }}</span></p>
                        <p><strong>中危漏洞:</strong> <span class="severity-medium">{{ report_data.statistics.medium }}</span></p>
                        <p><strong>低危漏洞:</strong> <span class="severity-low">{{ report_data.statistics.low }}</span></p>
                    </div>
                    
                    <h2>漏洞类型分布</h2>
                    <table>
                        <thead>
                            <tr>
                                <th>漏洞类型</th>
                                <th>数量</th>
                            </tr>
                        </thead>
                        <tbody>
                            {% for type, count in report_data.statistics.types.items() %}
                            <tr>
                                <td>{{ type }}</td>
                                <td>{{ count }}</td>
                            </tr>
                            {% endfor %}
                        </tbody>
                    </table>
                    
                    {% if report_data.vulnerabilities %}
                    <h2>漏洞详情</h2>
                    <table>
                        <thead>
                            <tr>
                                <th>漏洞ID</th>
                                <th>类型</th>
                                <th>路径</th>
                                <th>严重程度</th>
                                <th>风险评分</th>
                                <th>修复建议</th>
                            </tr>
                        </thead>
                        <tbody>
                            {% for vuln in report_data.vulnerabilities %}
                            <tr>
                                <td>{{ vuln.vuln_id }}</td>
                                <td>{{ vuln.vuln_type }}</td>
                                <td>{{ vuln.path }}</td>
                                <td>
                                    <span class="severity-{{ vuln.severity }}">{{ vuln.severity }}</span>
                                </td>
                                <td>{{ vuln.risk_score }}</td>
                                <td>{{ vuln.suggestion }}</td>
                            </tr>
                            {% endfor %}
                        </tbody>
                    </table>
                    {% endif %}
                </body>
                </html>
                ''')
            
            # 渲染模板
            html_content = template.render(report_data=report_data)
            return html_content
            
        except Exception as e:
            logger.error(f"生成HTML报告失败: {str(e)}")
            raise
    
    @staticmethod
    def _generate_pdf_report(report_data):
        """生成PDF格式的报告"""
        try:
            # 先生成HTML报告
            html_content = ReportService._generate_html_report(report_data)
            
            # 使用pdfkit将HTML转换为PDF
            # 注意：需要安装wkhtmltopdf并配置环境变量
            try:
                pdf_data = pdfkit.from_string(html_content, False)
                return pdf_data
            except OSError:
                # 如果没有安装wkhtmltopdf，返回HTML内容作为替代
                logger.warning("未找到wkhtmltopdf，无法生成PDF报告。请安装wkhtmltopdf后重试。")
                return ReportService._generate_html_report(report_data).encode('utf-8')
                
        except Exception as e:
            logger.error(f"生成PDF报告失败: {str(e)}")
            raise
    
    @staticmethod
    def _generate_excel_report(report_data):
        """生成Excel格式的报告"""
        try:
            # 创建Excel工作簿
            buffer = BytesIO()
            
            # 创建一个writer对象
            with pd.ExcelWriter(buffer, engine='xlsxwriter') as writer:
                # 扫描信息表
                scan_info_df = pd.DataFrame.from_dict([report_data['scan_info']])
                scan_info_df.to_excel(writer, sheet_name='扫描信息', index=False)
                
                # 统计信息表
                stats_df = pd.DataFrame.from_dict([report_data['statistics']])
                stats_df.to_excel(writer, sheet_name='统计信息', index=False)
                
                # 漏洞类型分布表
                types_df = pd.DataFrame(list(report_data['statistics']['types'].items()), 
                                       columns=['漏洞类型', '数量'])
                types_df.to_excel(writer, sheet_name='漏洞类型分布', index=False)
                
                # 漏洞详情表
                if report_data['vulnerabilities']:
                    vulns_df = pd.DataFrame(report_data['vulnerabilities'])
                    # 选择需要的列
                    columns_to_keep = ['vuln_id', 'vuln_type', 'path', 'severity', 'risk_score', 'suggestion']
                    vulns_df = vulns_df[[col for col in columns_to_keep if col in vulns_df.columns]]
                    vulns_df.to_excel(writer, sheet_name='漏洞详情', index=False)
                    
            # 获取Excel数据
            excel_data = buffer.getvalue()
            return excel_data
            
        except Exception as e:
            logger.error(f"生成Excel报告失败: {str(e)}")
            raise
    
    @staticmethod
    def export_vulnerabilities(scan_id, format_type='json'):
        """导出漏洞数据
        
        Args:
            scan_id (str): 扫描任务ID
            format_type (str): 导出格式，支持 'json', 'csv', 'excel'
            
        Returns:
            tuple: (导出数据, 文件扩展名, MIME类型)
        """
        try:
            # 获取漏洞数据
            vulnerabilities = Vulnerability.query.filter_by(scan_id=scan_id).all()
            if not vulnerabilities:
                raise ValueError(f"未找到ID为{scan_id}的扫描任务的漏洞数据")
            
            # 转换为字典列表
            vuln_list = [v.to_dict() for v in vulnerabilities]
            
            # 根据格式类型导出数据
            if format_type.lower() == 'json':
                return json.dumps(vuln_list, ensure_ascii=False, indent=2), 'json', 'application/json'
            elif format_type.lower() == 'csv':
                return ReportService._export_to_csv(vuln_list), 'csv', 'text/csv'
            elif format_type.lower() == 'excel':
                return ReportService._export_to_excel(vuln_list), 'xlsx', 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet'
            else:
                raise ValueError(f"不支持的导出格式: {format_type}")
                
        except Exception as e:
            logger.error(f"导出漏洞数据失败: {str(e)}")
            raise
    
    @staticmethod
    def _export_to_csv(vulnerabilities):
        """导出为CSV格式"""
        try:
            if not vulnerabilities:
                return ""
            
            # 创建DataFrame
            df = pd.DataFrame(vulnerabilities)
            
            # 选择需要的列
            columns_to_keep = ['vuln_id', 'vuln_type', 'path', 'severity', 'risk_score', 'suggestion']
            df = df[[col for col in columns_to_keep if col in df.columns]]
            
            # 转换为CSV字符串
            csv_data = df.to_csv(index=False, encoding='utf-8-sig')
            return csv_data
            
        except Exception as e:
            logger.error(f"导出CSV数据失败: {str(e)}")
            raise
    
    @staticmethod
    def _export_to_excel(vulnerabilities):
        """导出为Excel格式"""
        try:
            if not vulnerabilities:
                return b""
            
            # 创建Excel工作簿
            buffer = BytesIO()
            
            # 创建DataFrame
            df = pd.DataFrame(vulnerabilities)
            
            # 选择需要的列
            columns_to_keep = ['vuln_id', 'vuln_type', 'path', 'severity', 'risk_score', 'suggestion']
            df = df[[col for col in columns_to_keep if col in df.columns]]
            
            # 写入Excel
            with pd.ExcelWriter(buffer, engine='xlsxwriter') as writer:
                df.to_excel(writer, sheet_name='漏洞数据', index=False)
            
            # 获取Excel数据
            excel_data = buffer.getvalue()
            return excel_data
            
        except Exception as e:
            logger.error(f"导出Excel数据失败: {str(e)}")
            raise
    
    @staticmethod
    def get_report_templates():
        """获取可用的报告模板列表"""
        # 在实际应用中，这可能会从数据库或配置文件中读取模板列表
        templates = [
            {
                'id': 'default',
                'name': '默认报告模板',
                'description': '包含所有漏洞信息的标准报告格式'
            },
            {
                'id': 'summary',
                'name': '摘要报告模板',
                'description': '仅包含漏洞摘要和统计信息的简洁报告格式'
            },
            {
                'id': 'detailed',
                'name': '详细报告模板',
                'description': '包含完整漏洞详情、风险评估和修复建议的详细报告格式'
            }
        ]
        return templates